package com.liang.ali.yoloxgrpc1;

import com.google.protobuf.Int64Value;
import com.liang.ali.yoloxgrpc1.entity.FloatObject;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import net.coobird.thumbnailator.Thumbnails;
import org.tensorflow.framework.TensorProto;
import org.tensorflow.framework.TensorShapeProto;
import tensorflow.serving.Model;
import tensorflow.serving.Predict;
import tensorflow.serving.PredictionServiceGrpc;

import java.awt.image.BufferedImage;
import java.awt.image.Raster;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static com.liang.ali.yoloxgrpc1.PainImage.painImage;
import static org.tensorflow.framework.DataType.DT_FLOAT;


/*
 * @Author iLy
 * @Description
 * @ClassName test
 * @Date 2021/11/1 14:43
 * @VERSION 1.0
 **/
public class YoloxGrpc {

    public YoloxGrpc() {
    }

    public List<FloatObject> receiveCap(String fileName, String outputFolder) throws Exception {
        String outFolder = outputFolder;
        String file = fileName;
        //读取文件，强制修改图片大小，设置输出文件格式bmp(模型定义时输入数据是无编码的)
        BufferedImage im = Thumbnails.of(file).forceSize(640, 640).outputFormat("bmp").asBufferedImage();
        //转换图片到图片数组，匹配输入数据类型为Float
        Raster raster = im.getData();
        List<Float> floatList = new ArrayList<>();
        float[] temp = new float[raster.getWidth() * raster.getHeight() * raster.getNumBands()];
        float[] pixels = raster.getPixels(0, 0, raster.getWidth(), raster.getHeight(), temp);
        for (float pixel : pixels) {
            floatList.add(pixel/255);
        }
        for (int i = 0; i < floatList.size(); i++) {
            if (i % 3 == 0){
                Float aFloat = floatList.get(i);
                aFloat -= 0.485f;
                aFloat /= 0.229f;
                floatList.set(i, aFloat);
            }else if (i % 3 == 1){
                Float aFloat = floatList.get(i);
                aFloat -= 0.456f;
                aFloat /= 0.224f;
                floatList.set(i, aFloat);
            }else {
                Float aFloat = floatList.get(i);
                aFloat -= 0.406f;
                aFloat /= 0.225f;
                floatList.set(i, aFloat);
            }
        }

        List<Float> inputShape = new ArrayList<>();
        inputShape.add(416f);
        inputShape.add(416f);

        //#记个时
        long t = System.currentTimeMillis();
        //创建连接，注意usePlaintext设置为true表示用非SSL连接
        ManagedChannel channel = ManagedChannelBuilder.forAddress("121.196.146.141", 8506).usePlaintext(true).build();
        //这里还是先用block模式
        PredictionServiceGrpc.PredictionServiceBlockingStub stub = PredictionServiceGrpc.newBlockingStub(channel);
        //创建请求
        Predict.PredictRequest.Builder predictRequestBuilder = Predict.PredictRequest.newBuilder();
        //模型名称和模型方法名预设
        Model.ModelSpec.Builder modelSpecBuilder = Model.ModelSpec.newBuilder();
        modelSpecBuilder.setName("simple_test2");
        modelSpecBuilder.setSignatureName("serving_default");
        modelSpecBuilder.setVersion(Int64Value.of(1));
        predictRequestBuilder.setModelSpec(modelSpecBuilder);
        //设置入参,访问默认是最新版本，如果需要特定版本可以使用tensorProtoBuilder.setVersionNumber方法
        TensorProto.Builder tensorProtoBuilder = TensorProto.newBuilder();
        tensorProtoBuilder.setDtype(DT_FLOAT);
        TensorShapeProto.Builder tensorShapeBuilder = TensorShapeProto.newBuilder();
        tensorShapeBuilder.addDim(TensorShapeProto.Dim.newBuilder().setSize(1));
        //#150528 = 640 * 640 * 3
        tensorShapeBuilder.addDim(TensorShapeProto.Dim.newBuilder().setSize(640));
        tensorShapeBuilder.addDim(TensorShapeProto.Dim.newBuilder().setSize(640));
        tensorShapeBuilder.addDim(TensorShapeProto.Dim.newBuilder().setSize(3));
        tensorProtoBuilder.setTensorShape(tensorShapeBuilder.build());
        tensorProtoBuilder.addAllFloatVal(floatList);

        //设置入参,访问默认是最新版本，如果需要特定版本可以使用tensorProtoBuilder.setVersionNumber方法
        TensorProto.Builder tensorProtoBuilder1 = TensorProto.newBuilder();
        tensorProtoBuilder1.setDtype(DT_FLOAT);
        TensorShapeProto.Builder tensorShapeBuilder1 = TensorShapeProto.newBuilder();
        tensorShapeBuilder1.addDim(TensorShapeProto.Dim.newBuilder().setSize(1));
        //#150528 = 640 * 640 * 3
        tensorShapeBuilder1.addDim(TensorShapeProto.Dim.newBuilder().setSize(2));
        tensorProtoBuilder1.setTensorShape(tensorShapeBuilder1.build());
        tensorProtoBuilder1.addAllFloatVal(inputShape);

        predictRequestBuilder.putInputs("input_1", tensorProtoBuilder.build());
        predictRequestBuilder.putInputs("input_2", tensorProtoBuilder1.build());
        //访问并获取结果
        Predict.PredictResponse predictResponse = stub.predict(predictRequestBuilder.build());


        Map<String, TensorProto> outputsMap = predictResponse.getOutputsMap();
        List<Float> concatenate_13 = outputsMap.get("yolo_eval").getFloatValList();
        List<Float> concatenate_14 = outputsMap.get("yolo_eval_1").getFloatValList();
        List<Integer> concatenate_15 = outputsMap.get("yolo_eval_2").getIntValList();

        List<float[]> boxList = new ArrayList<>();
        int boxCount = concatenate_13.size() / 4;
        for (int i = 0; i < boxCount; i++) {
            float[] floats = new float[4];
            for (int i1 = 0; i1 < 4; i1++) {
                floats[i1] = concatenate_13.get(i * 4 + i1);
            }
            boxList.add(floats);
        }

        float[][] boxesResult = boxList.toArray(new float[boxList.size()][]);
        Float[] scoreResult = concatenate_14.toArray(new Float[concatenate_14.size()]);
        Integer[] classesResult = concatenate_15.toArray(new Integer[concatenate_15.size()]);

        painImage(file, classesResult, scoreResult, boxesResult, outFolder);

        System.out.println(concatenate_13);
        System.out.println(concatenate_14);
        System.out.println(concatenate_15);
        List<FloatObject> objects = new ArrayList<>();
        for (int i = 0; i < concatenate_15.size(); i++) {
            FloatObject floatObject = new FloatObject(concatenate_14.get(i), concatenate_15.get(i), boxList.get(i));
            objects.add(floatObject);
        }

        return objects;
    }
}



